import faiss
import numpy as np
import torch

from model.towers import ItemTower

class FaissIndexer:
    def __init__(self, config):
        self.config = config
        self.index = None
        self.item_ids = None
    
    def build_index(self, model):
        """构建Faiss索引"""
        # 加载物品塔
        item_tower = ItemTower(self.config.NUM_ITEMS, self.config.EMBEDDING_DIM)
        item_tower.load_state_dict(model.state_dict(), strict=False)
        item_tower.eval()
        
        # 生成所有物品嵌入
        all_items = torch.arange(self.config.NUM_ITEMS).long()
        with torch.no_grad():
            item_embeddings = item_tower(all_items).numpy()
        
        # 归一化
        faiss.normalize_L2(item_embeddings)
        
        # 创建索引
        if self.config.INDEX_TYPE == "FlatIP":
            self.index = faiss.IndexFlatIP(self.config.EMBEDDING_DIM)
        elif self.config.INDEX_TYPE == "IVFFlat":
            quantizer = faiss.IndexFlatIP(self.config.EMBEDDING_DIM)
            self.index = faiss.IndexIVFFlat(
                quantizer, 
                self.config.EMBEDDING_DIM, 
                100,  # nlist
                faiss.METRIC_INNER_PRODUCT
            )
            self.index.train(item_embeddings)
        elif self.config.INDEX_TYPE == "HNSW":
            self.index = faiss.IndexHNSWFlat(
                self.config.EMBEDDING_DIM, 
                32,  # M
                faiss.METRIC_INNER_PRODUCT
            )
            self.index.hnsw.efConstruction = 40
        
        # 添加向量到索引
        self.index.add(item_embeddings)
        self.item_ids = all_items.numpy()
        
        # 保存索引
        faiss.write_index(self.index, self.config.INDEX_PATH)
        np.save(self.config.INDEX_PATH.replace(".bin", "_ids.npy"), self.item_ids)
    
    def load_index(self):
        """加载预构建的索引"""
        self.index = faiss.read_index(self.config.INDEX_PATH)
        self.item_ids = np.load(self.config.INDEX_PATH.replace(".bin", "_ids.npy"))
    
    def search(self, user_emb, top_k):
        """搜索最相似的物品"""
        # 转换为numpy数组并归一化
        if isinstance(user_emb, torch.Tensor):
            user_emb = user_emb.detach().cpu().numpy()
        faiss.normalize_L2(user_emb)
        
        # 执行搜索
        distances, indices = self.index.search(user_emb, top_k)
        
        # 映射到原始物品ID
        item_ids = self.item_ids[indices]
        return item_ids, distances
